////////////////////////////////////
// TOPOLOGY CONTROL DAEMON
// (c) 2006 by Waltemar M. de Sousa Jr.
// waltemar@vt.edu

#define WIRELESS_EXTENSIONS
#define DEBUG
//#define DEBUG2

#include "tc_topology_controller.hpp"
#include "tc_db_manager.hpp"
#include "tc_msg_controller.hpp"
#include "tc_types.hpp"

#include <signal.h> // signal name macros, and the signal() prototype
#include <arpa/inet.h> // inet_addr() 
#include <sys/socket.h>
#include <netinet/in.h> // sockaddr_in, htonl(), htons()
//#include <net/if.h> // struct ifreq
//#include <unistd.h> // standard unix functions, e.g., alarm()
//#include <sys/ioctl.h> // SIOCGIFADDR

#if defined(WIRELESS_EXTENSIONS)
#include <iwlib.h> // wireless extensions API
#endif

#include <fstream>
#include <iostream>
#include <string>
#include <set>

using std::string;
using std::cout;
using std::endl;
using std::ifstream;
using std::set;


//////////////////////////////////
// ------------------ DEFINES ------------------ 

#define BUF_LEN 1500 // Maximum buffer length


//////////////////////////////////
// ---------- GLOBAL VARIABLES ----------- 

char MULTICAST_ADDR[20]; // Multicast group IP address
int  TC_PORT; // TC protocol port number
int  LOOPBACK; // Socket loopback option (0: off ; 1: on)
char ifname[10]; // wireless interface name
char iwpath[50]; // path to the WE tools
char IP_ADDR[20]; // local IP address
char MAC_ADDR[20] = "00:00:00:00:00:00"; // local MAC address
int  MAX_POWER_DBM; // maximum power level in dBm
int  TPC; // TPC capability [1: TPC capable]

#if defined(WIRELESS_EXTENSIONS)
struct ether_addr LOCAL_HWADDR; // local MAC address
#endif

int EXIT_ERROR_CODE = 0; // exit code issued by the SIGINT handler

struct sockaddr_in inAddr;  // Incoming address structure
struct sockaddr_in outAddr; // Outgoing address structure
int sid_server;             // Networking socket descriptor id
int skfd;                   // Generic raw socket descriptor to talk to the kernel

int TIME_COUNTER = -1; // Time variable used in the alarm handler
int MSG_COUNTER = 0;  // Specifies which message should be sent next
int ALARM_T = 1;      // Time interval (in sec) between alarm calls
int TX_INTERVAL = 2;  // Time interval (in sec) between transmissions
                      // Note about timers: for usec precision,
                      //                    use setitimer() instead of alarm()

int TC_TIMEOUT = 8;  // Time-out value (in sec) for the Topology Controller to reset
int TC_IDLE_T = 3;  // Time interval (in sec) between TC Run() calls in idle state
int TC_ROUTE_T = 2; // Time interval (in sec) between TC route info calls

int NODE_TIMEOUT = 5; //  Node time-out value in sec 
int EDGE_TIMEOUT = 5; //  Edge time-out value in sec 
int EDGE_WEIGHT_WINDOW = 4; // Max. number of elements in the weight moving average

int MAX_IW_POWER = 128; // maximum transmit power level allowed by iwpriv
int CUR_IW_POWER = 35;  // current transmit power level in the iwpriv scale

TCDatabaseManager tc_db;             // topology database module
TCDatabaseManager route_db;          // routing database module
TCTopologyController tc_control;     // topology controller module
TCMessageController tc_msg_ctrl;     // message controller module
TCStatsCollector tc_stats_collector; // statistics collector module

///////////////////////////////////
// --------- GLOBAL FUNCTIONS ------------

void catch_kill(int sig_num); // Prototype

///////////
// Print an error message and exit the program
void ExitOnError( string err )
{
  cout << err << endl;
  EXIT_ERROR_CODE = 1;
  catch_kill(SIGINT); // quit the program gracefully
}

///////////
// Send a P_MSG message on the socket with descriptor s
// Returns the number of bytes sent or -1 on error
// Note that size_t is equivalent to an unsigned long and that
// ssize_t is its signed version
ssize_t tc_p_msg_sendto(int s, p_msg *msg, int flags, 
			const struct sockaddr *to, socklen_t tolen)
{
  char buffer[BUF_LEN]; memset(&buffer, '0', BUF_LEN);
  ssize_t bytes_sent;
  size_t len = 0;
  
  char_hton( &buffer[len], &msg->TYPE, 2); len += 2;
  char_hton( &buffer[len], &msg->NODE.IP_ADDR, 4 ); len += 4;
  char_hton( &buffer[len], &msg->NODE.MAC_ADDR, 6); len += 6;
  char_hton( &buffer[len], &msg->NODE.TPC, 2 ); len += 2;
  char_hton( &buffer[len], &msg->NODE.MAX_POWER, 2 ); len += 2;

  bytes_sent = sendto(s, &buffer, len, flags, to, tolen);

  return bytes_sent;
}

///////////
// Send an R_MSG message on the socket with descriptor s
// Returns the number of bytes sent or -1 on error
ssize_t tc_r_msg_sendto(int s, r_msg *msg, int flags, 
			const struct sockaddr *to, socklen_t tolen)
{
  char buffer[BUF_LEN]; memset(&buffer, '0', BUF_LEN);
  ssize_t bytes_sent;
  size_t len = 0;
  int i;

  char_hton( &buffer[len], &msg->TYPE, 2); len += 2;
  char_hton( &buffer[len], &msg->LEN, 2); len += 2;
  
  for(i = 0; i < (int) msg->LEN; i++){
    if( (len + 16) > BUF_LEN )
      ExitOnError("R_MSG is larger than BUF_LEN.");
    char_hton( &buffer[len], &msg->NEIGHBOR[i].IP_ADDR, 4 ); len += 4;
    char_hton( &buffer[len], &msg->NEIGHBOR[i].MAC_ADDR, 6); len += 6;
    char_hton( &buffer[len], &msg->NEIGHBOR[i].TPC, 2 ); len += 2;
    char_hton( &buffer[len], &msg->NEIGHBOR[i].MAX_POWER, 2 ); len += 2;
    char_hton( &buffer[len], &msg->WEIGHT[i], 2 ); len += 2;
  }

  bytes_sent = sendto(s, &buffer, len, flags, to, tolen);

  return bytes_sent;
}

////////////
// Receive message on socket s (note that the msg must be decoded later)
ssize_t tc_recv(int s, void *buf, size_t len, int flags)
{
  ssize_t buflen;

  buflen = recv(s, buf, len, flags);

  return buflen;
}

/////////////
// Set TX power using the iwpriv command from the Wireless Extension Tools
//
//   NOTES: changing the power level with the hostap driver
//
//     iwpriv wlan0 alc 0 // disables auto-power-control
//     iwpriv wlan0 writemif 62 $LEVEL // sets txpower to $LEVEL
//     iwpriv wlan0 readmif 62 // reads register #(CR31)
//     iwpriv wlan0 readmif 116 // reads register #(CR58)
//     $LEVEL scale: max = 128 ; min = 127 ; middle = 0 or 255

void SetTxPower(int power_level) // note: power_level is in the driver scale
{
#if defined(WIRELESS_EXTENSIONS)
  static string path = (string) iwpath + "iwpriv";
  // iwpriv commands
  static string disable_auto_tpc = path + " " + ifname + " alc 0";
  string set_tx_power = path + " " + ifname +
                         " writemif 62 " + itos(power_level);

  // Always disable auto-TPC (do it again, just in case!)
  system(disable_auto_tpc.c_str());
  
  // Set the Tx power level
  system(set_tx_power.c_str());

#if defined(DEBUG)
  cout << endl << "SET TX POWER: " << set_tx_power << endl;
#endif

#endif // WIRELESS_EXTENSIONS
}


// -------------- SIGNAL HANDLERS ----------------

//////////
// Handler for SIGALRM
void catch_alarm(int sig_num){
  // Do we need to issue signal() again?
  // Answer: no, because we are already passing sig_num = SIGALRM as argument
  //signal(SIGALRM, catch_alarm);

  // Increment time counter by 1 sec
  ++TIME_COUNTER;

  // Update the timer in the topology database manager,
  //  routing database manager, and topology controller
  tc_db.UpdateTimer(TIME_COUNTER);
  route_db.UpdateTimer(TIME_COUNTER);
  tc_control.UpdateTimer(TIME_COUNTER);

#if defined(DEBUG)
  system("clear");
  cout << endl << "CLOCK = " << TIME_COUNTER << " sec" << endl;
  tc_db.PrintGraph();
//  cout << endl << "ROUTING TABLE" << endl;
//  route_db.PrintGraph();
#endif

  ///////////////////////
  // TC CONTROLLER EVENT
  if (TPC){
    tc_control.Run();
  }

  ///////////////////////
  // TRANSMISSION EVENTS
  if ((TIME_COUNTER % TX_INTERVAL) == 0){
    ssize_t buflen; // Counts the length of the transmitted buffer

    if(TPC){
      // --- SET the power level to MAX_IW_POWER ---
      //   note: power_level is in the driver scale
      SetTxPower(MAX_IW_POWER);
    }

    // -------- EVENT 0 ---------
    // ---    SEND A P_MSG    ---

    if (MSG_COUNTER == P_MSG){
      // Get the pMsg
      p_msg mypMsg;
      if (tc_msg_ctrl.AssembleMessage(mypMsg)){

        if( (buflen = tc_p_msg_sendto(sid_server, &mypMsg, 0,
				      (struct sockaddr *) &outAddr,
  				      sizeof(outAddr)) ) < 0 ){
  	  ExitOnError("tc_p_msg_sendto() failed.");
        }
    
#if defined(DEBUG)
        cout << endl << "Sending P_MSG... ";
        cout << buflen << " bytes sent." << endl;
#endif
      }
    }
    // --------- END EVENT 0 ---------

    // --------- EVENT 1 ---------
    // ---    SEND AN R_MSG    ---
    if (MSG_COUNTER == R_MSG){
      // Get the rMsg data.
      // First we need to find out the number of neighbors in order to create the
      // neighbor list and the edge weight list
      r_msg myrMsg;

      // Allocate memory for the neighbor list and weight list (variable length)
      int num_nodes = 1 + tc_db.GetNumberOfNeighbors(0);
      myrMsg.NEIGHBOR = (tc_node *) calloc(1, num_nodes * sizeof(tc_node));
      myrMsg.WEIGHT = (int16_t *) calloc(1, num_nodes * sizeof(int16_t));

      if (tc_msg_ctrl.AssembleMessage(myrMsg, num_nodes)){

        if( (buflen = tc_r_msg_sendto(sid_server, &myrMsg, 0,
				      (struct sockaddr *) &outAddr,
				      sizeof(outAddr)) ) < 0 ){
	  ExitOnError("tc_r_msg_sendto() failed.");
        }

#if defined(DEBUG)
      cout << endl << "Sending R_MSG... ";
      cout << buflen << " bytes sent." << endl;
#endif
      }

      // Deallocate memory for the neighbor list and the weight list
      free(myrMsg.NEIGHBOR);
      free(myrMsg.WEIGHT);
    }
    // --------- END EVENT 1 ---------

    if(TPC){
      // --- RESET the power level to CUR_IW_POWER ---
      //   note: power_level is in the driver scale
      SetTxPower(CUR_IW_POWER);
    }

  // Increment the message counter
  MSG_COUNTER = (MSG_COUNTER + 1) % 2;
  }

  // Set next alarm
  alarm(ALARM_T);
}

//////////
// Handler for SIGINT
void catch_kill(int sig_num)
{
  // Close the server socket
  if( sid_server ){
    if( close(sid_server) < 0 ){
      cout << "Error closing networking socket.";
      exit(1);
    }
  }

#if defined(WIRELESS_EXTENSIONS)
  if( skfd ){
    // Clear monitor list in the driver
    char buf[5] = "off";
    char** bufp = (char**) malloc(sizeof(char*));
    bufp[0] = &buf[0];
    if( set_spy_info(skfd, ifname, bufp, 1) < 0 ){
      cout << "set_spy_info() failed." << endl;
      exit(1);
    }
    free(bufp);

    // --- Re-enable automatic TPC ---
    if(TPC){
      string path = (string) iwpath + "iwpriv";
      string cmd = path + " " + ifname + " alc 1";
      system(cmd.c_str());

#if defined(DEBUG)
      cout << endl << "SET AUTO-TPC: " << cmd << endl;
#endif
    }

    // Close the raw socket
    iw_sockets_close(skfd);
  }
#endif // WIRELESS_EXTENSIONS

  cout << endl << "<<< Topology control daemon is now closed >>>" << endl << endl;
  exit(EXIT_ERROR_CODE);
}


////////////////////////////////////
// --------------------- MAIN( ) -----------------------
int main(int argc, char** argv)
{
  /////////////////////////////////////
  // --------------- SIGNAL HANDLERS -------------

  // Set a signal handler for ALRM signals
  signal(SIGALRM, catch_alarm);

  // Set a signal handler for KILL signals (ctrl-c)
  signal(SIGINT, catch_kill);


  ////////////////////////////////////////
  // --------- PARSE CONFIGURATION FILE ----------

  ifstream myfile ("TCd.config"); // Configuration file
  string line, lvalue, rvalue;
  string::size_type eqpos, pos;

  if (myfile.is_open()){ // Open configuration file
    while (!myfile.eof()){
      getline(myfile, line); // Read file line by line

      // Remove spaces
      while ( (pos = line.find(' ')) != string::npos ){
	line.erase(pos, 1);
      }

      // Process line
      if ( (line[0] == '#') || line.empty() ){
	continue;
      }

      // Find position of the '=' character
      if ( (eqpos = line.find('=')) == string::npos ){
	continue;
      }

      // Assign lvalue and rvalue
      lvalue = line.substr(0, eqpos);
      rvalue = line.substr(eqpos+1, line.size() - eqpos);

      if (lvalue == "MULTICAST_ADDR") strcpy(MULTICAST_ADDR, rvalue.c_str());
      if (lvalue == "TC_PORT") TC_PORT = atoi(rvalue.c_str());
      if (lvalue == "LOOPBACK") LOOPBACK = atoi(rvalue.c_str());
      if (lvalue == "INTERFACE") strcpy(ifname, rvalue.c_str());
      if (lvalue == "IP") strcpy(IP_ADDR, rvalue.c_str());
      if (lvalue == "MAC") strcpy(MAC_ADDR, rvalue.c_str());
      if (lvalue == "WE_PATH") strcpy(iwpath, rvalue.c_str());
      if (lvalue == "MAX_POWER_DBM") MAX_POWER_DBM = atoi(rvalue.c_str());
      if (lvalue == "TPC") TPC = (rvalue == "1" ? 1 : 0);
      if (lvalue == "TX_INTERVAL") TX_INTERVAL = atoi(rvalue.c_str());
      if (lvalue == "NODE_TIMEOUT") NODE_TIMEOUT = atoi(rvalue.c_str());
      if (lvalue == "EDGE_WEIGHT_WINDOW") EDGE_WEIGHT_WINDOW = atoi(rvalue.c_str());
      if (lvalue == "TC_TIMEOUT") TC_TIMEOUT = atoi(rvalue.c_str());
      if (lvalue == "TC_IDLE_T") TC_IDLE_T = atoi(rvalue.c_str());
      if (lvalue == "TC_ROUTE_T") TC_ROUTE_T = atoi(rvalue.c_str());
    }

    myfile.close(); // Close configuration file
  }
  else{
    cout << "Unable to open configuration file." << endl;
  }

#if defined(DEBUG)
  system("clear");
  cout << "Parsing configuration file:" << endl;
  cout << "---------------------------" << endl;
  cout << "MULTICAST_ADDR = " << MULTICAST_ADDR << endl;
  cout << "TC_PORT = " << TC_PORT << endl;
  cout << "LOOPBACK = " << (LOOPBACK ? "On" : "Off") << endl;
  cout << "IFNAME = " << ifname << endl;
  cout << "IP = " << IP_ADDR << endl;
  cout << "MAC = " << MAC_ADDR << endl;
  cout << "WE_PATH = " << iwpath << endl;
  cout << "MAX_POWER_DBM = " << MAX_POWER_DBM << endl;
  cout << "TPC = " << TPC << endl;
  cout << "TX_INTERVAL = " << TX_INTERVAL << endl;
  cout << "NODE_TIMEOUT = " << NODE_TIMEOUT << endl;
  cout << "EDGE_WEIGHT_WINDOW = " << EDGE_WEIGHT_WINDOW << endl;
  cout << "TC_TIMEOUT = " << TC_TIMEOUT << endl;
  cout << "TC_IDLE_T = " << TC_IDLE_T << endl;
  cout << "TC_ROUTE_T = " << TC_ROUTE_T << endl;

  cout << endl << "Press <ENTER> to continue..." << endl;
  getchar();
#endif


  /////////////////////////////////////
  // --------- SYSTEM INITIALIZATION  ----------

#if defined(WIRELESS_EXTENSIONS)
  // Create a channel to the NET kernel
  if( (skfd = iw_sockets_open()) < 0 )
    ExitOnError("Error opening raw socket.");

  // Get the MAC address if invalid
  unsigned short ptype;
  if (MAC_ADDR == "00:00:00:00:00:00"){
    iw_get_mac_addr(skfd, ifname, &LOCAL_HWADDR, &ptype);
  }
#endif

  // ------- TC DATABASE MANAGER --------

  // Set "use average" parameter (default is true)
  //tc_db.SetUseAvg( false );

  // Set topology database manager's NODE time-out value
  tc_db.SetTimeOut( (double) NODE_TIMEOUT );

  // Set topology database manager's EDGE time-out value
  EDGE_TIMEOUT = (EDGE_WEIGHT_WINDOW + 1) * TX_INTERVAL;
  tc_db.SetEdgeTimeOut( (double) EDGE_TIMEOUT );

  // Set topology database manager's EDGE WEIGHT WINDOW value
  tc_db.SetEdgeWindowSize( EDGE_WEIGHT_WINDOW );

  // Add self to the database manager
  tc_db.AddNode(IP_ADDR, MAC_ADDR, TPC, MAX_POWER_DBM);


  // ------- ROUTING DATABASE MANAGER --------

  // Set "use average" parameter (default is true)
  //route_db.SetUseAvg( false );

  // Set database manager's NODE time-out value
  route_db.SetTimeOut( (double) NODE_TIMEOUT );

  // Set database manager's EDGE time-out value
  route_db.SetEdgeTimeOut( (double) EDGE_TIMEOUT );

  // Set database manager's EDGE WEIGHT WINDOW value
  route_db.SetEdgeWindowSize( EDGE_WEIGHT_WINDOW );

  // Add self to the database manager
  route_db.AddNode(IP_ADDR, MAC_ADDR, TPC, MAX_POWER_DBM);


  // ------- TC TOPOLOGY CONTROLLER -------
  tc_control.Init(CUR_IW_POWER, tc_db, route_db, false,
                  TC_TIMEOUT, TC_IDLE_T, TC_ROUTE_T);


  // ------- STATISTICS COLLECTOR -------
  
  // Initialize the stats collector module
  tc_stats_collector.Init(skfd, ifname, tc_db);

  // ------- MESSAGE CONTROLLER --------

  // Initialize message controller module
  tc_msg_ctrl.Init(tc_db, tc_stats_collector);

  // ----- SOCKETS -----

  // Incoming address structure
  memset(&inAddr, 0, sizeof(inAddr)); // initialize saddr with zero
  inAddr.sin_family = AF_INET; // TCP/IP address family
  inAddr.sin_addr.s_addr = htonl( INADDR_ANY ); // Incoming address
  inAddr.sin_port = htons( TC_PORT ); // TC protocol port number

 // Open socket for network communications
  if( (sid_server = socket( PF_INET, SOCK_DGRAM, 0)) < 0 ){
    ExitOnError("Error opening networking socket.");
  }
  
  // Outgoing address structure 
  memset(&outAddr, 0, sizeof(outAddr)); // initialize saddr with zero
  outAddr.sin_family = AF_INET; // TCP/IP address family
  outAddr.sin_addr.s_addr = inet_addr( MULTICAST_ADDR ); // Outgoing address
  outAddr.sin_port = htons( TC_PORT ); // TC protocol port number


  /*  // Check the multicast interface (this is not working)
  struct in_addr addr; socklen_t size; size = sizeof(addr);
  if( getsockopt(sid_server, IPPROTO_IP, IP_MULTICAST_IF, &addr, &size) < 0 )
    ExitOnError("Error checking the multicast interface.");
  cout << endl << "Multicast interface address = " << addr.s_addr << endl;
  */


  // Enable/Disable the loopback option on the socket
  socklen_t size_loop = sizeof(LOOPBACK);
  if( setsockopt(sid_server, IPPROTO_IP,
		 IP_MULTICAST_LOOP, &LOOPBACK, size_loop) < 0 ){
    ExitOnError("Failed to disable the loopback option of the socket.");
  }


  // Join multicast group
  struct ip_mreq mreq;
  mreq.imr_multiaddr.s_addr=inet_addr( MULTICAST_ADDR );
  mreq.imr_interface.s_addr=htonl( INADDR_ANY );
  if (setsockopt(sid_server, IPPROTO_IP, IP_ADD_MEMBERSHIP, &mreq, sizeof(mreq)) < 0){
    ExitOnError("Failed to join multicast group.");
  }

  // Bind the socket to the incoming address
  if( bind(sid_server, (struct sockaddr *) &inAddr, sizeof(inAddr)) < 0 ){
    ExitOnError("Failed to bind the socket to the incoming address.");
  }

  
  ///////////////////////////////////////
  // ------- ENABLE TRANSMISSION LOOP ------

  // Call the alarm handler immediately
  //   - sets UNIX alarm for periodic transmission of messages;
  //   - alarm period = ALARM_T
  catch_alarm(SIGALRM);


  /////////////////////////////////////
  // --------------- RECEIVER LOOP ---------------
  
  char buf[BUF_LEN];
  ssize_t buflen; // Counts the length of the received buffer
  
  while(true){
    // Receive packets on the socket (this is a blocking I/O operation)
    buflen = 0;
    buflen = recv(sid_server, &buf, BUF_LEN, 0);


#if defined(DEBUG)
    cout << endl << "Received: " << buflen << " bytes." << endl;
#endif


    // Let the message controller handle msg parsing and interacting with other modules
    // Note: link statistics collection is done at GetStats() in the TCMessageController class
    tc_msg_ctrl.ParseMessage(buf, buflen);
  }

  // If, by changing the above code, we reach this point, then...
  // Quit application
  catch_kill(SIGINT);

  return 0;
}
